package com.structure;

import lombok.NoArgsConstructor;

/**
 * @Author: Gallrax
 * @Description: 非递归版本AVL树(非递归版本需要借助parentNode属性, 递归版本则不需要)
 * @Date: 2019/1/9
 */
public class AVLTree {

    private AVLNode root;

    private int height(AVLNode avlNode) {
        if (avlNode != null) {
            return avlNode.height;
        }
        return 0;
    }

    public AVLNode find(int key) {
        if (root == null || root.key == key) {
            return root;
        }
        AVLNode current = root;
        while (true) {
            if (current == null) {
                return null;
            }
            if (key < current.key) {
                current = current.leftNode;
            } else if (key > current.key) {
                current = current.rightNode;
            } else if (key == current.key) {
                return current;
            }
        }
    }

    public AVLNode insert(int key) {
        AVLNode avlNode = new AVLNode(key);
        if (root == null) {
            avlNode.height = 1;
            root = avlNode;
            return avlNode;
        }
        AVLNode parent = root;
        while (true) {
            if (key < parent.key) {
                if (parent.leftNode != null) {
                    parent = parent.leftNode;
                } else {
                    parent.leftNode = avlNode;
                    avlNode.parentNode = parent;
                    //非递归式从底层往上层检查平衡
                    while (parent != null) {
                        //因为是往左插,则肯定可能是左侧比右侧高
                        if (height(parent.leftNode) - height(parent.rightNode) == 2) {
                            if (key < parent.leftNode.key) {
                                //如果插入的key是小于parent左侧，则一定是左左插入，因此左左旋转
                                parent = leftLeftRotation(parent);
                            } else {
                                //如果插入的key不是小于parent左侧，则一定是左右插入，因此左右旋转
                                parent = leftRightRotation(parent);
                            }
                        }
                        parent = parent.parentNode;
                        parent.height = Math.max(height(parent.leftNode), height(parent.rightNode)) + 1;
                    }
                }
            } else {
                if (parent.rightNode != null) {
                    parent = parent.rightNode;
                } else {
                    parent.rightNode = avlNode;
                    avlNode.parentNode = parent;
                    //非递归式从底层往上层检查平衡
                    while (parent != null) {
                        //因为是往右插,则肯定可能是右侧比左侧高
                        if (height(parent.rightNode) - height(parent.leftNode) == 2) {
                            if (key > parent.rightNode.key) {
                                //如果插入的key是大于parent右侧，则一定是右右插入，因此右右旋转
                                parent = leftLeftRotation(parent);
                            } else {
                                //如果插入的key不是大于parent左侧，则一定是右左插入，因此右左旋转
                                parent = leftRightRotation(parent);
                            }
                        }
                        parent = parent.parentNode;
                    }
                }
            }
        }
    }

    public boolean delete(int key) {
        if (root != null) {
            return false;
        }
        AVLNode current = root;
        AVLNode parent = current;
        while (true) {
            parent = current;
            if (current == null) {
                return false;
            } else if (key < current.key) {
                if (current.leftNode != null) {
                    current = current.leftNode;
                } else {
                    return false;
                }
            } else if (key > current.key) {
                if (current.rightNode != null) {
                    current = current.rightNode;
                } else {
                    return false;
                }
            } else {
                //如果删除节点没有子节点,则只需将父节点对其引用删除即可
                int childCount = 0;
                if (current.leftNode != null) {
                    childCount++;
                }
                if (current.rightNode != null) {
                    childCount++;
                }
                switch (childCount) {
                    case 0:
                        return deleteNoneChild(parent, current);
                    case 1:
                        return deleteOneChild(parent, current);
                    case 2:
                        return deleteTwoChild(parent, current);
                    default:
                        return false;
                }
            }
        }
    }

    private AVLNode leftLeftRotation(AVLNode avlNode) {
        //旋转后的根节点
        AVLNode current = avlNode.leftNode;
        //1.1将旋转后的根节点的右子节点赋予给avlNode的左子节点
        avlNode.leftNode = current.rightNode;
        //1.2将avlNode赋予给旋转后的根节点右子节点的父节点
        current.rightNode.parentNode = avlNode;
        //2将avlNode赋予给旋转后的根节点的右子节点
        current.rightNode = avlNode;
        //3.1将avlNode的父节点赋予给旋转后的根节点的父节点
        current.parentNode = avlNode.parentNode;
        //3.2将旋转后的根节点点赋予给avlNode的父节点
        avlNode.parentNode = current;
        //设置高度
        avlNode.height = Math.max(height(avlNode.leftNode), height(avlNode.rightNode)) + 1;
        current.height = Math.max(height(current.leftNode), height(current.rightNode)) + 1;
        return current;
    }

    private AVLNode leftRightRotation(AVLNode avlNode) {
        //先右右旋转
        avlNode.leftNode = rightRightRotation(avlNode.leftNode);
        //再进行左左旋转
        return leftLeftRotation(avlNode);
    }

    private AVLNode rightRightRotation(AVLNode avlNode) {
        //旋转后的根节点
        AVLNode current = avlNode.rightNode;
        //1.1将旋转后的根节点的左子节点赋予给avlNode的右子节点
        avlNode.rightNode = current.leftNode;
        //1.2将avlNode赋予给旋转后的根节点的左子节点的父节点
        current.leftNode.parentNode = avlNode;
        //2将avlNode赋予给旋转后的根节点的左子节点
        current.leftNode = avlNode;
        //3.1将avlNode的父节点赋予给旋转后的根节点的父节点
        current.parentNode = avlNode.parentNode;
        //3.2将旋转后的根节点赋予给avlNode的父节点
        avlNode.parentNode = current;
        //设置高度
        avlNode.height = Math.max(height(avlNode.leftNode), height(avlNode.rightNode)) + 1;
        current.height = Math.max(height(current.leftNode), height(current.rightNode)) + 1;
        return current;
    }

    private AVLNode rightLeftRotation(AVLNode avlNode) {
        avlNode.rightNode = leftLeftRotation(avlNode.rightNode);
        return rightRightRotation(avlNode);
    }

    private boolean deleteNoneChild(AVLNode parent, AVLNode avlNode) {
        if (avlNode == parent.leftNode) {
            parent.leftNode = null;
        } else {
            parent.rightNode = null;
        }
        return true;
    }

    private boolean deleteOneChild(AVLNode parent, AVLNode avlNode) {
        if (avlNode == parent.leftNode) {
            parent.leftNode = parent.leftNode != null ? parent.leftNode : parent.rightNode;
        } else {
            parent.rightNode = parent.leftNode != null ? parent.leftNode : parent.rightNode;
        }
        return true;
    }

    private boolean deleteTwoChild(AVLNode parent, AVLNode avlNode) {
        AVLNode successor = findSuccessor(avlNode);
        if (avlNode == parent.leftNode) {
            parent.leftNode = successor;
        } else {
            parent.rightNode = successor;
        }
        return true;
    }

    private AVLNode findSuccessor(AVLNode avlNode) {
        AVLNode current = avlNode.rightNode;
        AVLNode parent = current;
        //如果删除节点的右子节点没有左子节点,则删除节点的右子节点即为后继节点
        if (current.leftNode == null) {
            return current;
        }
        while (current != null) {
            parent = current;
            current = current.leftNode;
        }
        parent.leftNode = current.rightNode;
        current.rightNode = avlNode.rightNode;
        return current;
    }


}

@NoArgsConstructor
class AVLNode {
    int key;
    int height;
    AVLNode leftNode;
    AVLNode rightNode;
    AVLNode parentNode;

    public AVLNode(int key) {
        this.key = key;
    }
}
